# IMPORT LIBRERIE
from ultralytics import YOLO
import cv2
import os

# CONFIGURATION
MODEL_PATH = "C:/Users/caruso/Desktop/YOLO_Project/agrilus/train/weights/best.pt"  # Select YOLO model
TEST_IMAGE_NAME = "PRI0000251_laterodorsal.png"  # Example of image in the dataset
TEST_IMAGE_PATH = f"C:/Users/caruso/Desktop/YOLO_Project/images/test_unknown/{TEST_IMAGE_NAME}"  # Image path

CONF_THRESHOLD = 0.3

# 1. LOAD THE MODEL
print("YOLO Template Upload...")
model = YOLO(MODEL_PATH)

# 2. CHECK IMAGE
if not os.path.exists(TEST_IMAGE_PATH):
    print(f"Error: Image {TEST_IMAGE_PATH} does not exist. Check the path and extension.")
    exit()

# 3. UPLOAD AND VIEW IMAGE
print(f"Loading test image: {TEST_IMAGE_NAME}")
image = cv2.imread(TEST_IMAGE_PATH)

# 4. MAKE PREDICTION
print("Making Prediction...")
results = model.predict(image, conf=CONF_THRESHOLD)

# 5. VIEW RESULTS
print("\n**Prediction results**")
for result in results:
    if result.names is not None and len(result.names) > 0:
        print(f"Predicted Class: {result.names[0]} with {result.probs.top1conf*100:.2f}% confidence")
    else:
        print("No classes detected!")

# 6. SAVE IMAGE
output_path = f"C:/Users/caruso/Desktop/YOLO_Project/images/test_unknown/{TEST_IMAGE_NAME}"
os.makedirs(os.path.dirname(output_path), exist_ok=True)
result_image = results[0].plot()
cv2.imwrite(output_path, result_image)
print(f"Image with results saved in: {output_path}")
